# 实现几何形貌相关的特征，包括法向量场，曲率场
import numpy as np
from .contact import get_contact_region

def get_normal_field(position: np.ndarray) -> np.ndarray:
    """
    计算法向量场
    Args:
        position (np.ndarray): 点云位置，形状为 (400, 3) 或 (20, 20, 3) 或 (N, 400, 3) 或 (N, 20, 20, 3)
    Returns:
        np.ndarray: 法向量场，形状为(20, 20, 3) 或 (N, 20, 20, 3)
    """
    if position.shape[1:] != (400, 3) and position.shape[1:] != (20, 20, 3) and position.shape != (20, 20, 3) and position.shape != (400, 3):
        raise ValueError("Position must be of shape (N, 400, 3) or (N, 20, 20, 3) or (400,3) or (20,20,3)")
    if position.shape == (400, 3):
        position = position.reshape((20, 20, 3))
    if position.shape[1:] == (400,3):
        position = position.reshape((-1, 20, 20, 3))

    normals = np.zeros_like(position)
    x_diff = np.diff(position, axis=-3)
    y_diff = np.diff(position, axis=-2)
    shape: tuple = position.shape
    zero_shape_x = (shape[0], 1, shape[2], shape[3]) if position.ndim == 4 else (1, shape[1], shape[2])
    zero_shape_y = (shape[0], shape[1], 1, shape[3]) if position.ndim == 4 else (shape[0], 1, shape[2])
    diff_x1 = np.concatenate((x_diff, np.zeros(zero_shape_x)), axis=-3) # p[i+1, j] - p[i, j]
    diff_x2 = np.concatenate((np.zeros(zero_shape_x), -x_diff), axis=-3) # p[i-1, j] - p[i, j]
    diff_y1 = np.concatenate((y_diff, np.zeros(zero_shape_y)), axis=-2) # p[i, j+1] - p[i, j]
    diff_y2 = np.concatenate((np.zeros(zero_shape_y), -y_diff), axis=-2) # p[i, j-1] - p[i, j]

    # cross product to get normals
    normals = np.cross(diff_x1, diff_y1) + np.cross(diff_y1,diff_x2) + np.cross(diff_x2,diff_y2) + np.cross(diff_y2,diff_x1)

    # normalize normals
    norms = np.linalg.norm(normals, axis= -1, keepdims=True)
    normals = normals / (norms + 1e-6) # (N, 20, 20, 3) or (20, 20, 3)

    # adjust direction to make z component positive
    neg_mask = normals[..., 2] < 0
    normals[neg_mask] = -normals[neg_mask]
    return normals.reshape(-1,400,3).squeeze()

def get_curvature_field(position: np.ndarray, normal: np.ndarray) -> np.ndarray:
    raise NotImplementedError("Curvature field computation is not implemented yet.")

def get_average_normal(P: np.ndarray, F: np.ndarray) -> np.ndarray:
    '''
    Get the average normal vector of the 3D Points, weighted by the magnitude of the 3D Forces.
    Args:
        P (np.ndarray): 3D Points of shape (N, 400, 3).
        F (np.ndarray): 3D Forces of shape (N, 400, 3).
    Returns:
        np.ndarray: The average normal vector of shape (N, 3).
    '''
    assert P.shape == F.shape and P.shape[-1] == 3, "P and F must have the same shape and last dimension must be 3."
    assert P.shape[1:] == (400,3) or P.shape[1:] == (20,20,3), "P and F must have shape (N, 400, 3) or (N, 20,20, 3)."
    normal = get_normal_field(P)  # (N, 400, 3)
    P = P.reshape(P.shape[0], -1, 3)  # (N, 400, 3)
    F = F.reshape(F.shape[0], -1, 3)  # (N, 400, 3)
    # weighted by force magnitude
    weight = get_contact_region(P, F, method="soft_thres", threshold=0.01)  # (N, 400)
    weighted_normal = normal * weight[..., np.newaxis]  # (N, 400, 3)
    average_normal = weighted_normal.sum(axis=-2) / weight.sum(axis=-1)[..., np.newaxis]  # (N, 3)
    average_normal = average_normal / np.linalg.norm(average_normal, axis=-1)[..., np.newaxis]  # (N, 3)
    return average_normal